//! Trait declaration parser

use super::super::{parse_type, Parser};
use crate::ast::decl::*;
use crate::ast::types::Type;
use crate::error::{CompilerError, Result};
use crate::lexer::token::Token;
use super::helpers::parse_where_clause_for_generics;

/// Parse trait declaration (type Name trait { ... })
pub fn parse_trait_declaration(parser: &mut Parser) -> Result<TraitDecl> {
    parser.expect(&Token::Type)?;

    // Parse trait name
    let trait_name = if let Some(Token::Identifier(name)) = parser.peek().map(|t| &t.token) {
        let name = name.clone();
        parser.advance()?;
        name
    } else {
        return Err(CompilerError::syntax(
            parser.current_location().line,
            parser.current_location().column,
            "Expected trait name after 'type'",
        ));
    };

    // Parse generic type parameters (they come after identifier, before 'trait' keyword)
    let mut generics = if parser.consume(&Token::Less)? {
        let mut generics = Vec::new();

        if !parser.check(&Token::Greater) {
            loop {
                if let Some(Token::Identifier(generic_name)) = parser.peek().map(|t| &t.token) {
                    let name = generic_name.clone();
                    parser.advance()?;

                    // Parse constraints
                    let constraints = if parser.consume(&Token::Colon)? {
                        let mut constraints = Vec::new();

                        loop {
                            if let Some(constraint) = parse_type(parser)? {
                                constraints.push(constraint);
                            }

                            if !parser.consume(&Token::Plus)? {
                                break;
                            }
                        }

                        constraints
                    } else {
                        vec![]
                    };

                    generics.push(GenericParam {
                        name,
                        constraints,
                        location: parser.current_location(),
                    });
                } else {
                    break;
                }

                if !parser.consume(&Token::Comma)? {
                    break;
                }
            }
        }

        parser.expect(&Token::Greater)?;
        generics
    } else {
        vec![]
    };

    // Expect 'trait' keyword
    parser.expect(&Token::Trait)?;

    // Optional where-clause between `trait` and `{`
    parse_where_clause_for_generics(parser, &mut generics)?;

    // Parse opening brace
    parser.expect(&Token::LeftBrace)?;

    // Parse trait methods
    let mut methods = Vec::new();

    while !parser.check(&Token::RightBrace) {
        if parser.is_at_end() {
            return Err(CompilerError::syntax(
                parser.current_location().line,
                parser.current_location().column,
                "Expected '}' to close trait declaration",
            ));
        }

        if let Some(method) = parse_trait_method(parser)? {
            methods.push(method);
        } else {
            // Skip unknown tokens
            parser.advance()?;
        }
    }

    parser.expect(&Token::RightBrace)?;

    Ok(TraitDecl {
        name: trait_name,
        generics,
        methods,
        location: parser.current_location(),
    })
}

/// Parse trait declaration from current position (assumes 'type' and name are already consumed)
/// This is used when we've already peeked ahead to determine it's a trait
/// Generic parameters may be present after the identifier, before 'trait' keyword
pub fn parse_trait_declaration_from_current(
    parser: &mut Parser,
    trait_name: String,
) -> Result<TraitDecl> {
    // Parse generic type parameters (they come after identifier, before 'trait' keyword)
    let mut generics = if parser.consume(&Token::Less)? {
        let mut generics = Vec::new();

        if !parser.check(&Token::Greater) {
            loop {
                if let Some(Token::Identifier(generic_name)) = parser.peek().map(|t| &t.token) {
                    let name = generic_name.clone();
                    parser.advance()?;

                    // Parse constraints
                    let constraints = if parser.consume(&Token::Colon)? {
                        let mut constraints = Vec::new();

                        loop {
                            if let Some(constraint) = parse_type(parser)? {
                                constraints.push(constraint);
                            }

                            if !parser.consume(&Token::Plus)? {
                                break;
                            }
                        }

                        constraints
                    } else {
                        vec![]
                    };

                    generics.push(GenericParam {
                        name,
                        constraints,
                        location: parser.current_location(),
                    });
                } else {
                    break;
                }

                if !parser.consume(&Token::Comma)? {
                    break;
                }
            }
        }

        parser.expect(&Token::Greater)?;
        generics
    } else {
        vec![]
    };

    // Now expect 'trait' keyword
    parser.expect(&Token::Trait)?;

    // Optional where-clause between `trait` and `{`
    parse_where_clause_for_generics(parser, &mut generics)?;

    // Parse opening brace
    parser.expect(&Token::LeftBrace)?;

    // Parse trait methods
    let mut methods = Vec::new();

    while !parser.check(&Token::RightBrace) {
        if parser.is_at_end() {
            return Err(CompilerError::syntax(
                parser.current_location().line,
                parser.current_location().column,
                "Expected '}' to close trait declaration",
            ));
        }

        if let Some(method) = parse_trait_method(parser)? {
            methods.push(method);
        } else {
            // Skip unknown tokens
            parser.advance()?;
        }
    }

    parser.expect(&Token::RightBrace)?;

    Ok(TraitDecl {
        name: trait_name,
        generics,
        methods,
        location: parser.current_location(),
    })
}

/// Parse a trait method signature
pub(in crate::parser) fn parse_trait_method(parser: &mut Parser) -> Result<Option<TraitMethod>> {
    if !parser.check(&Token::Fn) {
        return Ok(None);
    }

    parser.expect(&Token::Fn)?;

    // Parse method name
    let method_name = if let Some(Token::Identifier(name)) = parser.peek().map(|t| &t.token) {
        let name = name.clone();
        parser.advance()?;
        name
    } else {
        return Err(CompilerError::syntax(
            parser.current_location().line,
            parser.current_location().column,
            "Expected method name after 'func'",
        ));
    };

    // Parse parameters
    parser.expect(&Token::LeftParen)?;
    let mut parameters = Vec::new();

    if !parser.check(&Token::RightParen) {
        loop {
            // Check for &self parameter first
            if parser.check(&Token::Ampersand) {
                parser.advance()?; // consume '&'
                if let Some(Token::Identifier(param_name)) = parser.peek().map(|t| &t.token) {
                    if param_name == "self" {
                        parser.advance()?; // consume 'self'

                        // For trait methods, &self doesn't have an explicit type
                        // We'll use a placeholder type that will be resolved later
                        let param_type = Type::Reference(crate::ast::types::ReferenceType {
                            referenced_type: Box::new(Type::Basic(
                                crate::ast::types::BasicType::Any,
                            )),
                            location: parser.current_location(),
                        });

                        parameters.push(crate::ast::types::Parameter {
                            name: "self".to_string(),
                            param_type,
                            default_value: None,
                            location: parser.current_location(),
                        });

                        // Check if there are more parameters or if this is the end
                        if parser.check(&Token::Comma) {
                            parser.advance()?; // consume comma, continue to next parameter
                            continue;
                        } else if parser.check(&Token::RightParen) {
                            // No more parameters, break out of loop
                            break;
                        } else {
                            // Unexpected token
                            return Err(CompilerError::syntax(
                                parser.current_location().line,
                                parser.current_location().column,
                                "Expected ',' or ')' after &self parameter",
                            ));
                        }
                    }
                }
                return Err(CompilerError::syntax(
                    parser.current_location().line,
                    parser.current_location().column,
                    "Expected 'self' after '&' in parameter",
                ));
            }

            // Parse parameter name first (support both NRC syntax: name type and traditional: name: type)
            if let Some(Token::Identifier(param_name)) = parser.peek().map(|t| &t.token) {
                let name = param_name.clone();
                parser.advance()?;

                // Check if there's a colon (traditional syntax)
                let param_type = if parser.consume(&Token::Colon)? {
                    // Traditional syntax: name: type
                    let ty = parse_type(parser)?;
                    if ty.is_none() {
                        return Err(CompilerError::syntax(
                            parser.current_location().line,
                            parser.current_location().column,
                            "Expected parameter type after colon",
                        ));
                    }
                    ty.unwrap()
                } else {
                    // NRC/Go-style syntax: name type
                    let ty = parse_type(parser)?;
                    if ty.is_none() {
                        return Err(CompilerError::syntax(
                            parser.current_location().line,
                            parser.current_location().column,
                            "Expected parameter type",
                        ));
                    }
                    ty.unwrap()
                };

                parameters.push(crate::ast::types::Parameter {
                    name,
                    param_type,
                    default_value: None,
                    location: parser.current_location(),
                });
            } else {
                break;
            }

            if !parser.consume(&Token::Comma)? {
                break;
            }
        }
    }

    parser.expect(&Token::RightParen)?;

    // Parse return type (support both arrow syntax and direct syntax like Go)
    let return_type = if parser.consume(&Token::Arrow)? {
        // Arrow syntax: func name() -> Type;
        let ty = parse_type(parser)?;
        parser.expect(&Token::Semicolon)?; // Semicolon required after arrow syntax
        ty
    } else if parser.check(&Token::Semicolon) {
        // No return type: func name();
        parser.expect(&Token::Semicolon)?; // Consume the semicolon
        None
    } else {
        // For direct syntax without arrow, check if next token is a type
        // First check if this could be direct syntax by looking ahead
        if let Some(Token::Identifier(_)) = parser.peek().map(|t| &t.token) {
            // Save current position in case this isn't a type
            let current_pos = parser.current.clone();

            if let Ok(Some(ty)) = parse_type(parser) {
                // Successfully parsed a type - check what comes after
                if parser.check(&Token::Semicolon) {
                    // This is direct syntax: func name() Type;
                    parser.expect(&Token::Semicolon)?; // consume the semicolon
                    Some(ty)
                } else if parser.check(&Token::RightBrace) {
                    // This is direct syntax without semicolon: func name() Type }
                    Some(ty)
                } else {
                    // This doesn't match expected patterns, restore position
                    parser.current = current_pos;
                    None
                }
            } else {
                // Not a valid type, restore position
                parser.current = current_pos;
                None
            }
        } else {
            // No potential type found
            None
        }
    };

    Ok(Some(TraitMethod {
        name: method_name,
        parameters,
        return_type,
        location: parser.current_location(),
    }))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::parser::Parser;

    fn parse_trait_decl(input: &str) -> Result<TraitDecl> {
        let mut parser = Parser::new(input.to_string(), None);
        // Initialize parser by advancing to first token
        parser.advance()?;
        parse_trait_declaration(&mut parser)
    }

    #[test]
    fn test_parse_simple_trait() {
        // Use actual syntax with &self parameter (matches actual file syntax)
        let result = parse_trait_decl("type Display trait { func to_string(&self) String }");
        if let Err(e) = &result {
            eprintln!("Parse error: {:?}", e);
        }
        assert!(result.is_ok(), "Failed to parse trait: {:?}", result);
        let decl = result.unwrap();
        assert_eq!(decl.name, "Display");
        assert_eq!(decl.methods.len(), 1);
        assert_eq!(decl.methods[0].name, "to_string");
    }

    #[test]
    fn test_parse_trait_with_multiple_methods() {
        // Use actual syntax with &self parameters (matches actual file syntax)
        let result = parse_trait_decl(
            "type Drawable trait { func draw(&self); func get_position(&self) Point }",
        );
        if let Err(e) = &result {
            eprintln!("Parse error: {:?}", e);
        }
        assert!(
            result.is_ok(),
            "Failed to parse trait with multiple methods: {:?}",
            result
        );
        let decl = result.unwrap();
        assert_eq!(decl.name, "Drawable");
        assert_eq!(decl.methods.len(), 2);
    }

    #[test]
    fn test_parse_trait_with_generics() {
        // Use direct syntax without semicolon
        let result = parse_trait_decl("type Container<T> trait { func get(index: int) T }");
        match &result {
            Ok(decl) => {
                assert_eq!(decl.name, "Container");
                assert_eq!(decl.generics.len(), 1);
                assert_eq!(decl.generics[0].name, "T");
            }
            Err(e) => {
                panic!("Failed to parse trait with generics: {:?}", e);
            }
        }
    }

    #[test]
    fn test_parse_trait_empty() {
        let result = parse_trait_decl("type Empty trait {}");
        if let Err(e) = &result {
            eprintln!("Parse error: {:?}", e);
        }
        assert!(result.is_ok(), "Failed to parse empty trait: {:?}", result);
        let decl = result.unwrap();
        assert_eq!(decl.name, "Empty");
        assert_eq!(decl.methods.len(), 0);
    }

    #[test]
    fn test_parse_trait_method_with_parameters() {
        // Use direct syntax without semicolon
        let result = parse_trait_decl("type Calculator trait { func add(a: int, b: int) int }");
        if let Err(e) = &result {
            eprintln!("Parse error: {:?}", e);
        }
        assert!(
            result.is_ok(),
            "Failed to parse trait method with parameters: {:?}",
            result
        );
        let decl = result.unwrap();
        assert_eq!(decl.methods[0].parameters.len(), 2);
    }
}
